from pytorch_lightning.callbacks import ModelCheckpoint,Callback
import time
import wandb
from datam import *
from model import Model,Backbone
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning import Trainer
from config import  config
from callback import LogPredictionSamplesCallback


if __name__ == '__main__':


    pl.seed_everything(1234)
    # log model only if `val_accuracy` increases
    wandb_logger = WandbLogger(log_model="all", config=config.__dict__,
                               name=time.strftime("%Y_%m_%d_%H_%M_%S", time.localtime(time.time())))
    checkpoint_callback = [LogPredictionSamplesCallback(),ModelCheckpoint(monitor="val_accuracy", mode="max")]
    trainer = Trainer(gpus=config.AVAIL_GPUS, max_epochs=config.max_epochs, callbacks=checkpoint_callback,logger=wandb_logger)
    data_mnist = DataM(config.data_dir,config.BATCH_SIZE,config.AVAIL_GPUS)

    net = Backbone()
    model = Model(net)

    #训练模型
    trainer.fit(model,data_mnist)
    trainer.save_checkpoint(config.ckpt_path)
    trainer.test(model,data_mnist)





